from .base_model import BaseModel
import os.path as osp

import torch
import numpy as np
import cv2
from .FaceBoxes.FaceBoxes import FaceBoxes,viz_bbox
# import time


class FaceBoxModel(BaseModel):
    """
    This class implements the CycleGAN model, for learning image-to-image translation without paired data.

    The model training requires '--dataset_mode unaligned' dataset.
    By default, it uses a '--netG resnet_9blocks' ResNet generator,
    a '--netD basic' discriminator (PatchGAN introduced by pix2pix),
    and a least-square GANs objective ('--gan_mode lsgan').

    CycleGAN paper: https://arxiv.org/pdf/1703.10593.pdf
    """
    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        """Add new dataset-specific options, and rewrite default values for existing options.

        Parameters:
            parser          -- original option parser
            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.

        Returns:
            the modified parser.

        For CycleGAN, in addition to GAN losses, we introduce lambda_A, lambda_B, and lambda_identity for the following losses.
        A (source domain), B (target domain).
        Generators: G_A: A -> B; G_B: B -> A.
        Discriminators: D_A: G_A(A) vs. B; D_B: G_B(B) vs. A.
        Forward cycle loss:  lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper)
        Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper)
        Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 "Photo generation from paintings" in the paper)
        Dropout is not used in the original CycleGAN paper.
        """
        parser.set_defaults(no_dropout=True)  # default CycleGAN did not use dropout
        # parser.set_defaults(frequency=True)
        parser.add_argument('--inputs', type=str, default='img', help='the moudal of input data ')
        # parser.add_argument('--num_classes', type=int, default=3, help='the architecture situation of model')
        # parser.add_argument('--pretrained_model', type=str, default="lightcnn", help='the architecture of pretrained model')
        # parser.add_argument('--pretrained_path', type=str, default="lightcnn", help='the path of pretrained model')
        # parser.add_argument('--pretrained_ndf', type=int, default=128, help='the number of features of output ')
        # parser.add_argument('--pretrained_output_nc', type=int, default=256, help='the number of features of output ')

        return parser

    def __init__(self, opt):
        """Initialize the CycleGAN class.

        Parameters:
            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        
        BaseModel.__init__(self, opt)
        self.face_boxes = FaceBoxes(timer_flag=False)

    
    def set_input(self, input):
        """Unpack input data from the dataloader and perform necessary pre-processing steps.

        Parameters:
            input (dict): include the data itself and its metadata information.

        The option 'direction' can be used to swap domain A and domain B.
        """
        self.img = input[self.opt.inputs]
        pass
    def forward(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
        # self.fake_B = self.netG_A(self.real_A)  # G_A(A)
        # self.rec_A = self.netG_B(self.fake_B)   # G_B(G_A(A))
        # self.fake_A = self.netG_B(self.real_B)  # G_B(B)
        # self.rec_B = self.netG_A(self.fake_A)   # G_A(G_B(B))
        
        dets = self.face_boxes(self.img)  # xmin, ymin, w, h
        # print(dets)

        # repeating inference for `n` times
        # n = 10
        # for i in range(n):
        #     dets = self.face_boxes(self.img)
        
        self.results = viz_bbox(self.img, dets)


    def optimize_parameters(self):
        """Calculate losses, gradients, and update network weights; called in every training iteration"""
        pass
    def iter_end(self):
        if self.isTrain:
                pass
        elif self.opt.controller != 'test':
            pass
        elif self.opt.controller == 'test':
            pass
        
    def epoch_end(self):

        if self.isTrain:
            pass
        elif self.opt.controller != 'test':
            pass
        elif self.opt.controller == 'test':
            pass
    
    def display_opt_init(opt):
        # parser.add_argument('--inputs', type=str, default='img', help='the moudal of input data ')
        opt.__setattr__('inputs','img')
        return opt
    def display_init(self,_dict):
        # parser.add_argument('--inputs', type=str, default='img', help='the moudal of input data ')
        return _dict
    def display_proc(self,_dict):
        self.results['img'] = self.img
        return self.results